from openopt import NLP, GLP
from numpy import cos, arange, ones, asarray, abs, zeros, sqrt, asscalar, array, pi
from pylab import legend, show, plot, subplot, xlabel, subplots_adjust
from string import rjust, ljust, expandtabs

import JT2 as JT
import Calibration


#===============================================================================
#
#    Data Paths
#
#===============================================================================

fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
#fixedImageName = 'C:/Users/bryan/bryan-code/trunk/Images/CalibratedDRRImage.mhd'
inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
calFileLat = 'C:/Users/bryan/bryan-code/trunk/TestData/ext_cal1.txt'
roi = [[210,187],[277,370]]

#===============================================================================
#    Setup DRR
#===============================================================================

xraycam = JT.XRayCamera()
drr = JT.DRR()
drr.SetXRayCamera(xraycam)
drr.SetBlackOnWhiteImage()
drr.InteractionOff()

cal = Calibration.ExternalCalibration()
cal.LoadConsolidatedCalibration(calFileLat)
cal.LoadStaFile(staFile)

xraycam.SetExternalCalibration(cal)

volume=JT.Volume()
volume.SetVolumeFilename(inputVolFileName)
volume.SetOriginalTransform(cal._VolumeT)
volume.SetBodyCenter((3.906, 2.17, -7.378))
drr.AddVolume(volume)

# Set Fixed image using path
fixedImage = JT.FixedImage()
fixedImage.SetFileName(fixedImageName)

reg=JT.NewRegistration()
reg.SetFixedImage(fixedImage.GetImage(0))
reg.SetMovingImageGenerator(drr)

reg.SetRegionOfInterest(roi)

#===============================================================================
#    Setup Optimizer
#===============================================================================
method = 'ncc'

reg.SetImageMetric(method)
probName = method.upper()+'\n'+fixedImageName[fixedImageName.rfind('/')+1:]

nlp_list=['ralg', 'scipy_cobyla', 'lincher', 'scipy_slsqp', 'ipopt','algencan']
glp_list=['pswarm','galileo','de']

#solver = nlp_list[2]
solver = glp_list[1]

#For Testing
T = JT.Testing()
pose = T.GetStartPose(50)
print "Starting Pose: ", pose
fnGoal = 'minimize'

#lowerBounds=array([2,2,2,1,1,5])
search_range = 1*array([1,1,1,pi/180,pi/180,pi/180])
lowerBounds = pose - search_range
upperBounds = pose + search_range
maxTime=120
maxTimeCPU=120
plotFlag=1
x0 = pose
f=reg.GetValue

beg_steps = 1#[.1,.1,.1,.1,.1,.1]
end_steps = .001#[.1,.1,.1,.01,.01,.01]

#x0=array((0,0,0,0,0,0))
#print "Image tweaked: ", x0
#drr.GenerateImage(x0)

#
if solver == 'scipy_cobyla':
    print "Starting guess: ", x0
    p = NLP(f, x0, lb=lowerBounds,  ub=upperBounds,
        maxIter = 200,  maxFunEvals = 200, plot = plotFlag,
        rhobeg = beg_steps, rhoend = end_steps, iprint = 3,
        maxTime = maxTime,  maxCPUTime = maxTimeCPU, name=probName)
elif solver in nlp_list:
    print "Starting guess: ", x0
    p = NLP(f, x0, lb=lowerBounds,  ub=upperBounds,
        maxIter = 1000,  maxFunEvals = 500, plot = plotFlag,
        maxTime = maxTime,  maxCPUTime = maxTimeCPU, name=probName)
elif solver in glp_list:
    p = GLP(f, lb=lowerBounds,  ub=upperBounds,
        maxIter = 50,  maxFunEvals = 1000, plot = plotFlag,
        maxTime = maxTime,  maxCPUTime = maxTimeCPU, name=probName)
else:
    print "Not familiar with the how to handle %s solver"%solver

#optional: graphic output
#p.plot = 1

r = p.solve(solver, plot=plotFlag, debug=0)
x_opt,  f_opt = r.xf,  r.ff


print "Results:"
print "    optimum value: ", f_opt
print "    optimum param: ", x_opt
print "Starting Pose    : ", pose

